{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e8d99f0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchaudio\n",
    "def sisnr(x, s, eps=1e-8):\n",
    "    \"\"\"\n",
    "    calculate training loss\n",
    "    input:\n",
    "          x: separated signal, N x S tensor\n",
    "          s: reference signal, N x S tensor\n",
    "    Return:\n",
    "          sisnr: N tensor\n",
    "    \"\"\"\n",
    "\n",
    "    def l2norm(mat, keepdim=False):\n",
    "        return torch.norm(mat, dim=-1, keepdim=keepdim)\n",
    "\n",
    "    if x.shape != s.shape:\n",
    "        raise RuntimeError(\n",
    "            \"Dimention mismatch when calculate si-snr, {} vs {}\".format(\n",
    "                x.shape, s.shape))\n",
    "    x_zm = x - torch.mean(x, dim=-1, keepdim=True)\n",
    "    s_zm = s - torch.mean(s, dim=-1, keepdim=True)\n",
    "    t = torch.sum(\n",
    "        x_zm * s_zm, dim=-1,\n",
    "        keepdim=True) * s_zm / (l2norm(s_zm, keepdim=True)**2 + eps)\n",
    "    return 20 * torch.log10(eps + l2norm(t) / (l2norm(x_zm - t) + eps))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6739ea75",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([200, 200])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "x = torch.full((2,), 200)\n",
    "print(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "521c12e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.size(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aecc5576",
   "metadata": {},
   "outputs": [],
   "source": [
    "mix_audio, _ = torchaudio.load(\"/Netdata/2021/zb/data/LibriMix/Libri2Mix/wav16k/min/dev/mix_clean/1272-128104-0000_2035-147961-0014.wav\")\n",
    "clean_audio, _ = torchaudio.load(\"/Netdata/2021/zb/data/LibriMix/Libri2Mix/wav16k/min/dev/s1/1272-128104-0000_2035-147961-0014.wav\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "fba5de46",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([2.1145])\n"
     ]
    }
   ],
   "source": [
    "print(sisnr(mix_audio,clean_audio))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a34b1d70",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sisnr(mix_audio,mix_audio))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "c56a74e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['epoch10.pth', 'epoch22.pth']"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## Test RE\n",
    "import re\n",
    "files = [\"epoch10.pth\", \"epoch22.pth\", \"epoch33.pth\", \"best.pth\", \"epoch_50.pth\"]\n",
    "import re\n",
    "\n",
    "epoch = 24  # Your threshold\n",
    "\n",
    "def _find_files(files, epoch, res = r\"epoch(\\d+)\\.pth$\"):\n",
    "    \"\"\"\n",
    "    Find files that is smaller than the epoch according to a specific regular expression.\n",
    "    Returns:\n",
    "        Return the files that is smaller than epoch\n",
    "    \"\"\"\n",
    "    # Filter files that match the pattern and extract the number\n",
    "    matched = []\n",
    "    for f in files:\n",
    "        m = re.search(res, f)\n",
    "        if m:\n",
    "            if int(m.group(1)) < epoch:\n",
    "                matched.append((int(m.group(1)), f))  # (epoch_number, filename)\n",
    "    # Sort by epoch number\n",
    "    matched.sort()\n",
    "    # Extract only filenames (if needed)\n",
    "    sorted_files = [f for _, f in matched]\n",
    "    return sorted_files\n",
    "\n",
    "\n",
    "files = _find_files(files, epoch)\n",
    "files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "69ae4f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bd1b1a1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.random.rand(200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "984f14aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(400,)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a = np.tile(x, 2)\n",
    "a.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0df8edcc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "200"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c061155",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test-if librispeech has the same spks as libri2mix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f3e67edd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from pathlib import Path\n",
    "from typing import Dict, List, Union\n",
    "\n",
    "def read_2column_text(path: Union[Path, str]) -> Dict[str, str]:\n",
    "    \"\"\"Read a text file having 2 column as dict object.\n",
    "\n",
    "    Examples:\n",
    "        wav.scp:\n",
    "            key1 /some/path/a.wav\n",
    "            key2 /some/path/b.wav\n",
    "\n",
    "        >>> read_2column_text('wav.scp')\n",
    "        {'key1': '/some/path/a.wav', 'key2': '/some/path/b.wav'}\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    data = {}\n",
    "    with Path(path).open(\"r\", encoding=\"utf-8\") as f:\n",
    "        for linenum, line in enumerate(f, 1):\n",
    "            sps = line.rstrip().split(maxsplit=1)\n",
    "            if len(sps) == 1:\n",
    "                k, v = sps[0], \"\"\n",
    "            else:\n",
    "                k, v = sps\n",
    "            if k in data:\n",
    "                raise RuntimeError(f\"{k} is duplicated ({path}:{linenum})\")\n",
    "            data[k] = v\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1443689e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1172\n",
      "{'118', '1502', '3914', '5519', '1740', '8684', '1482', '2674', '2775', '3781', '561', '2272', '4854', '1355', '4267', '548', '3446', '114', '1195', '6038', '451', '8848', '4260', '4681', '850', '7828', '549', '7445', '925', '4848', '8791', '1334', '1069', '39', '6359', '4586', '8419', '5147', '340', '7188', '8421', '2581', '2473', '8300', '3537', '9023', '1841', '8324', '55', '198', '2416', '8410', '6147', '1624', '5400', '1060', '298', '1052', '1425', '1806', '8113', '3374', '7520', '7128', '922', '3830', '7078', '8095', '7754', '1316', '2688', '7732', '229', '7090', '6865', '1112', '1058', '3851', '8118', '2929', '7783', '8825', '4071', '2256', '8592', '208', '3728', '6904', '543', '3294', '7705', '4859', '5022', '7794', '3157', '699', '1923', '7957', '2618', '5867', '2823', '3513', '3584', '6574', '227', '1390', '2045', '2146', '8195', '7594', '835', '3526', '8605', '9022', '5727', '3328', '7730', '7398', '8152', '8388', '7120', '6308', '345', '6690', '2238', '4973', '7484', '4064', '1046', '3389', '475', '6258', '979', '1182', '1311', '5039', '559', '8051', '7069', '7126', '7276', '3717', '1826', '7095', '5570', '6373', '1903', '3521', '445', '1859', '6458', '1212', '1851', '2093', '3105', '3242', '4160', '8725', '226', '6081', '7285', '2388', '1649', '606', '7294', '3370', '3889', '8008', '1289', '6538', '1335', '205', '4438', '5514', '1382', '5049', '38', '6836', '1031', '1313', '2741', '7241', '8011', '7867', '8699', '4290', '2240', '954', '2254', '5703', '7264', '497', '5740', '5538', '1054', '8722', '2289', '274', '3112', '716', '196', '83', '4837', '5984', '5767', '7312', '1343', '1241', '311', '1603', '1754', '2110', '7495', '8028', '5154', '7030', '32', '22', '7316', '5093', '480', '8887', '2229', '98', '4425', '612', '7178', '203', '1553', '1224', '1283', '1460', '7402', '2691', '2598', '166', '669', '337', '5940', '6782', '6981', '70', '3094', '830', '7752', '1066', '6339', '6235', '2334', '93', '1825', '5968', '5054', '1322', '6763', '2182', '6032', '5337', '7169', '730', '4427', '625', '2533', '78', '6104', '1509', '7286', '23', '1222', '5809', '7780', '5724', '434', '175', '7434', '1034', '1246', '3852', '7000', '587', '3630', '2774', '6115', '6937', '1547', '7247', '8824', '1536', '7437', '56', '882', '698', '7766', '8464', '249', '4434', '2393', '7720', '8677', '3699', '4830', '246', '816', '4226', '339', '5126', '8142', '1001', '6531', '1365', '5007', '8479', '4238', '2893', '500', '7117', '6563', '696', '1571', '2992', '5139', '101', '6317', '6643', '1226', '8225', '472', '3224', '4051', '953', '594', '7258', '4340', '3448', '3436', '8630', '7478', '8222', '302', '1165', '2319', '7335', '948', '1018', '3072', '7825', '3540', '6848', '5063', '4770', '2853', '4133', '6206', '5776', '7314', '4362', '4039', '7555', '576', '7910', '6300', '4018', '7859', '8138', '5157', '8014', '8228', '6098', '6575', '4246', '3235', '8347', '1556', '3259', '6064', '380', '5909', '2577', '289', '4733', '7569', '986', '1513', '6788', '3879', '5731', '2758', '2004', '1678', '3082', '6544', '1053', '525', '1789', '8772', '3977', '7011', '7302', '4214', '8108', '887', '2299', '4289', '4629', '5239', '5622', '597', '5333', '454', '90', '54', '8119', '3546', '1116', '4381', '7318', '2531', '6918', '3982', '5750', '4744', '8494', '3119', '7505', '2815', '3230', '5802', '7837', '6529', '8629', '3118', '7962', '3092', '6492', '7383', '7416', '150', '6286', '7938', '8975', '1638', '125', '7789', '5339', '2787', '224', '8855', '6406', '6437', '2971', '6014', '1535', '4110', '5561', '5242', '5660', '868', '3792', '2910', '7498', '288', '3114', '7802', '446', '8401', '7517', '5810', '5456', '8747', '724', '7800', '5489', '6476', '8404', '8797', '329', '8573', '3869', '4856', '8183', '3857', '1898', '254', '8050', '1811', '580', '3905', '1992', '8329', '8425', '8580', '1446', '7832', '1455', '2167', '1337', '6294', '4719', '2532', '3083', '7460', '8687', '6395', '412', '8465', '1012', '4481', '3274', '7777', '7926', '3357', '5868', '6269', '6157', '2384', '1849', '2427', '225', '4297', '6215', '3003', '6371', '6925', '8459', '7190', '8468', '201', '5192', '362', '7384', '4195', '1777', '6189', '7816', '459', '4057', '3607', '1933', '1175', '1705', '3703', '8545', '5448', '7134', '2816', '8879', '89', '1235', '258', '2039', '2843', '1160', '6446', '688', '8820', '3046', '4839', '3228', '5588', '7635', '464', '7447', '81', '6367', '1874', '4335', '781', '1401', '5678', '4111', '2156', '7833', '6818', '7717', '1473', '2989', '7739', '806', '2911', '6696', '2230', '6895', '2517', '3972', '5883', '4138', '2592', '3180', '4957', '5401', '1264', '2007', '2696', '6519', '492', '3922', '7226', '14', '1417', '4598', '8238', '6341', '664', '5513', '639', '8006', '4807', '1578', '6828', '1447', '4860', '6426', '8176', '4010', '458', '4441', '1387', '7991', '2494', '671', '1422', '1743', '2137', '7647', '7949', '2514', '984', '831', '19', '3549', '4734', '8786', '3486', '8506', '5712', '6637', '8534', '3258', '8498', '1987', '1265', '3638', '4967', '6454', '6550', '7518', '4640', '2570', '7981', '3723', '4257', '6120', '1913', '2196', '7525', '2053', '5133', '1779', '2391', '7558', '6078', '4088', '6494', '115', '815', '1263', '2085', '5778', '3551', '278', '6080', '1885', '6378', '783', '476', '7956', '4490', '233', '403', '16', '2194', '6060', '192', '5914', '122', '1413', '3490', '6352', '6388', '8635', '4098', '3240', '1734', '6233', '5123', '6288', '7367', '6075', '7229', '8080', '1498', '87', '8057', '2562', '1383', '3347', '4243', '159', '6673', '3664', '3289', '4356', '5618', '1629', '1271', '2481', '2589', '323', '1970', '3330', '210', '2812', '2512', '5012', '1081', '5104', '335', '7511', '589', '6956', '4945', '4013', '3923', '3967', '5062', '8194', '8643', '2092', '126', '4137', '3994', '2582', '487', '8075', '4595', '441', '5206', '6509', '2397', '2060', '7868', '353', '6119', '1336', '242', '5637', '834', '4278', '250', '6927', '7051', '8591', '4152', '8875', '1121', '3866', '7059', '6694', '2436', '803', '8123', '5463', '1552', '2654', '7085', '176', '1974', '8190', '481', '820', '512', '296', '3008', '5092', '154', '1183', '157', '2285', '359', '4898', '2002', '6385', '398', '40', '3947', '4495', '8718', '1061', '2769', '2162', '8226', '100', '2960', '7939', '7140', '5115', '5293', '3825', '3835', '4145', '7733', '8163', '3307', '7994', '6415', '119', '2518', '3185', '836', '4535', '163', '7538', '2628', '369', '5604', '6555', '1594', '2159', '231', '4406', '207', '8770', '79', '770', '328', '1379', '7481', '2498', '1323', '3864', '511', '1845', '4680', '911', '2827', '636', '3686', '3733', '4899', '1088', '4236', '3876', '5389', '4397', '1607', '5583', '1958', '2751', '5304', '7395', '6686', '2074', '204', '1040', '1634', '1348', '5386', '3440', '4731', '8713', '216', '4363', '1027', '8063', '6139', '6683', '8088', '6924', '667', '1645', '708', '6167', '8490', '2952', '4592', '28', '4116', '7339', '2573', '1731', '3025', '5723', '1463', '1914', '1827', '6727', '405', '2638', '1445', '1724', '4358', '7342', '7933', '3483', '1769', '27', '3379', '3380', '4813', '2499', '510', '7113', '2790', '7297', '217', '6499', '663', '1926', '272', '1639', '5652', '3340', '637', '2149', '5002', '5684', '7139', '2709', '7540', '8609', '1098', '1050', '1448', '439', '534', '7874', '4014', '8396', '8705', '6877', '7240', '1641', '1800', '2294', '7688', '9026', '1723', '303', '3615', '2136', '6082', '5808', '1867', '4806', '968', '3983', '7959', '6965', '460', '374', '4926', '5190', '6497', '17', '408', '8758', '839', '5319', '3493', '8575', '7313', '583', '1079', '949', '6701', '1363', '2920', '8098', '3927', '6567', '5606', '7909', '3790', '3945', '7553', '7278', '8097', '1093', '581', '60', '7067', '1349', '2113', '3032', '4800', '2673', '1748', '1737', '3989', '7475', '2201', '5876', '7665', '4788', '7881', '3368', '2348', '2817', '8838', '7657', '7945', '3168', '7995', '2368', '248', '200', '2836', '2652', '30', '373', '3738', '5186', '1100', '4054', '3187', '5290', '409', '6518', '1472', '3215', '501', '5985', '1752', '3816', '829', '5746', '3979', '1487', '5390', '7704', '2401', '5163', '4433', '8066', '909', '1801', '3221', '6505', '6620', '211', '596', '64', '711', '7809', '2204', '4846', '2411', '3482', '3645', '731', '5688', '5672', '7061', '103', '426', '1259', '4044', '2364', '6000', '4148', '1392', '6006', '4853', '1668', '593', '6553', '5029', '8771', '8312', '6037', '5635', '764', '8193', '5655', '6188', '6019', '6272', '2404', '6510', '6099', '3654', '209', '5261', '332', '6209', '1296', '5918', '920', '1776', '957', '7148', '7245', '5266', '2882', '3009', '6181', '707', '8776', '3361', '2127', '5975', '1028', '4222', '8742', '5789', '1290', '6054', '598', '3001', '7932', '188', '5656', '899', '5717', '5935', '6330', '1456', '2012', '3070', '5393', '4519', '6993', '2056', '7967', '8527', '1025', '718', '3214', '3807', '8474', '6160', '322', '7982', '8266', '307', '240', '5246', '7145', '666', '2999', '2061', '1963', '6880', '1961', '4590', '5189', '479', '7515', '1944', '2269', '2764', '5322', '1943', '26', '318', '8619', '4331', '112', '2785', '3171', '2010', '4108'}\n"
     ]
    }
   ],
   "source": [
    "libri2mix_scp = read_2column_text(\"/Netdata/2021/zb/data/LibriMix/Libri2Mix/wav16k/min/lists/train/all/s1.scp\")\n",
    "spks = set([i.split(\"-\")[0] for i in list(libri2mix_scp.keys())])\n",
    "print(len(spks))\n",
    "print(spks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b39c39e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1172\n"
     ]
    }
   ],
   "source": [
    "with open(\"/DKUdata/tangbl/speech-seperation/tse/spex_plus/data/librispeech/spks.txt\", \"r\") as f:\n",
    "    l = len(f.readlines())\n",
    "print(l)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "317a50a2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'4108\\n'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "res[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "646449b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "40\n",
      "{'1988', '7850', '2078', '777', '3536', '6313', '1919', '2428', '652', '8842', '3752', '3000', '3576', '1462', '5536', '6241', '6319', '7976', '5694', '3081', '3170', '1272', '2035', '422', '2086', '2277', '1993', '6295', '1673', '5338', '8297', '2902', '2803', '174', '6345', '5895', '3853', '2412', '251', '84'}\n"
     ]
    }
   ],
   "source": [
    "libri2mix_scp = read_2column_text(\"/Netdata/2021/zb/data/LibriMix/Libri2Mix/wav16k/min/lists/dev/s1.scp\")\n",
    "spks = set([i.split(\"-\")[0] for i in list(libri2mix_scp.keys())])\n",
    "print(len(spks))\n",
    "print(spks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "581021e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/DKUdata/tangbl/data/librispeech/train_clean_100_360.pkl\", 'rb') as f:\n",
    "    spk_dict = pickle.load(f)\n",
    "\n",
    "spks_librispeech = set(list(spk_dict.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "19f8ced3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1172\n"
     ]
    }
   ],
   "source": [
    "print(len(spks_librispeech))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "bdfc21f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spks_librispeech == spks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "74f2d479",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/DKUdata/tangbl/speech-seperation/tse/spex_plus/data/librispeech/spks.txt\" ,\"w\") as f:\n",
    "    for _spk in spks_librispeech:\n",
    "        f.write(f\"{_spk}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "73f6a834",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bbd61c2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.arange(0,10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7bcee223",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    }
   ],
   "source": [
    "if (x == 1).any():\n",
    "    print(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c3237e65",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(5)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y  = torch.tensor(2)\n",
    "y + 3"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bltang_py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
